from flask import Flask, request, jsonify, render_template, redirect, url_for
import sqlite3
import json
import uuid
from datetime import datetime
import qrcode
from io import BytesIO
import base64

app = Flask(__name__)

# 初始化数据库
def init_db():
    conn = sqlite3.connect('votes.db')
    c = conn.cursor()
    
    # 创建投票表
    c.execute('''CREATE TABLE IF NOT EXISTS polls
                 (id TEXT PRIMARY KEY,
                  title TEXT NOT NULL,
                  options TEXT NOT NULL,
                  multiple_choice BOOLEAN DEFAULT 0,
                  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
    
    # 创建投票记录表
    c.execute('''CREATE TABLE IF NOT EXISTS votes
                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
                  poll_id TEXT NOT NULL,
                  option_index INTEGER NOT NULL,
                  voter_id TEXT NOT NULL,
                  voted_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                  FOREIGN KEY (poll_id) REFERENCES polls (id))''')
    
    conn.commit()
    conn.close()

# 首页
@app.route('/')
def index():
    return render_template('index.html')

# 创建投票
@app.route('/create', methods=['POST'])
def create_poll():
    data = request.json
    poll_id = str(uuid.uuid4())[:8]
    
    conn = sqlite3.connect('votes.db')
    c = conn.cursor()
    c.execute('''INSERT INTO polls (id, title, options, multiple_choice)
                 VALUES (?, ?, ?, ?)''',
              (poll_id, data['title'], json.dumps(data['options']), data['multiple_choice']))
    conn.commit()
    conn.close()
    
    return jsonify({'poll_id': poll_id, 'short_url': f'/poll/{poll_id}'})

# 获取投票页面
@app.route('/poll/<poll_id>')
def poll_page(poll_id):
    return render_template('poll.html', poll_id=poll_id)

# 获取投票详情
@app.route('/api/poll/<poll_id>')
def get_poll(poll_id):
    conn = sqlite3.connect('votes.db')
    c = conn.cursor()
    
    # 获取投票信息
    c.execute('SELECT title, options, multiple_choice FROM polls WHERE id = ?', (poll_id,))
    poll = c.fetchone()
    
    if not poll:
        return jsonify({'error': '投票不存在'}), 404
    
    title, options, multiple_choice = poll
    options = json.loads(options)
    
    # 获取投票统计
    c.execute('SELECT option_index, COUNT(*) FROM votes WHERE poll_id = ? GROUP BY option_index', (poll_id,))
    votes = c.fetchall()
    
    vote_counts = [0] * len(options)
    for option_index, count in votes:
        if 0 <= option_index < len(options):
            vote_counts[option_index] = count
    
    total_votes = sum(vote_counts)
    percentages = [0] * len(options)
    if total_votes > 0:
        percentages = [(count / total_votes) * 100 for count in vote_counts]
    
    conn.close()
    
    return jsonify({
        'title': title,
        'options': options,
        'multiple_choice': bool(multiple_choice),
        'vote_counts': vote_counts,
        'percentages': percentages,
        'total_votes': total_votes
    })

# 投票
@app.route('/api/vote/<poll_id>', methods=['POST'])
def vote(poll_id):
    data = request.json
    voter_id = str(uuid.uuid4())[:8]
    
    conn = sqlite3.connect('votes.db')
    c = conn.cursor()
    
    # 检查投票是否存在
    c.execute('SELECT multiple_choice FROM polls WHERE id = ?', (poll_id,))
    poll = c.fetchone()
    
    if not poll:
        conn.close()
        return jsonify({'error': '投票不存在'}), 404
    
    multiple_choice = bool(poll[0])
    
    # 检查是否已经投过票
    c.execute('SELECT COUNT(*) FROM votes WHERE poll_id = ? AND voter_id = ?', (poll_id, voter_id))
    if c.fetchone()[0] > 0 and not multiple_choice:
        conn.close()
        return jsonify({'error': '您已经投过票了'}), 400
    
    # 记录投票
    for option_index in data['options']:
        c.execute('INSERT INTO votes (poll_id, option_index, voter_id) VALUES (?, ?, ?)',
                  (poll_id, option_index, voter_id))
    
    conn.commit()
    conn.close()
    
    return jsonify({'success': True})

# 生成二维码
@app.route('/api/qrcode/<poll_id>')
def generate_qr(poll_id):
    url = f"{request.host_url}poll/{poll_id}"
    
    qr = qrcode.QRCode(version=1, box_size=10, border=5)
    qr.add_data(url)
    qr.make(fit=True)
    
    img = qr.make_image(fill_color="black", back_color="white")
    
    buffer = BytesIO()
    img.save(buffer, format='PNG')
    buffer.seek(0)
    
    img_str = base64.b64encode(buffer.getvalue()).decode()
    
    return jsonify({'qr_code': f"data:image/png;base64,{img_str}"})

# 获取投票详情（管理员）
@app.route('/api/admin/<poll_id>')
def get_poll_admin(poll_id):
    conn = sqlite3.connect('votes.db')
    c = conn.cursor()
    
    # 获取投票信息
    c.execute('SELECT title, options, multiple_choice FROM polls WHERE id = ?', (poll_id,))
    poll = c.fetchone()
    
    if not poll:
        conn.close()
        return jsonify({'error': '投票不存在'}), 404
    
    title, options, multiple_choice = poll
    options = json.loads(options)
    
    # 获取详细投票记录
    c.execute('''SELECT v.option_index, v.voter_id, v.voted_at
                 FROM votes v
                 WHERE v.poll_id = ?
                 ORDER BY v.voted_at DESC''', (poll_id,))
    votes = c.fetchall()
    
    detailed_votes = []
    for option_index, voter_id, voted_at in votes:
        if 0 <= option_index < len(options):
            detailed_votes.append({
                'option': options[option_index],
                'voter_id': voter_id,
                'voted_at': voted_at
            })
    
    conn.close()
    
    return jsonify({
        'title': title,
        'options': options,
        'multiple_choice': bool(multiple_choice),
        'detailed_votes': detailed_votes,
        'total_votes': len(detailed_votes)
    })

if __name__ == '__main__':
    init_db()
    app.run(debug=True, host='0.0.0.0', port=5000)
